import os
import paddlehub as hub
import cv2
import math
import numpy as np
import tqdm


## 定义realsr模型
sr_model = hub.Module(name='realsr')
##
# 加载模型
# 可根据需要更换上述模型中的任何一个
module_name = 'falsr_c'

import glob
import os

paths = []
def walk_type(path, file_type):
    paths = glob.glob(os.path.join(path,#存放图片的文件夹路径
                                   file_type # 文件类型，*代表所有文件
                                   )
                      )# path下所有file_type类型的文件的路径列表
    return paths
path_y = walk_type('/mnt/data/paddelocr/pubtabnet/pubtabnet/val','*.png')
print(path_y)
len(path_y)


def ChaoFen(paths):
    cropSize = 250
    img = cv2.imread(paths)
    hh, ww = img.shape[0:2]

    # 图片尺寸若大于切分尺寸则取imgList图片进行超分辨率处理
    if hh > cropSize and ww > cropSize:
        hCutNum = math.ceil(hh / cropSize)
        wCutNum = math.ceil(ww / cropSize)
        imgList = []

        ##循环先行后列，入队列则会是先列后行，与后续拼接时的顺序对应回来

        for hnum in range(hCutNum):
            for wnum in range(wCutNum):
                sub = img[hnum * cropSize:min((hnum + 1) * cropSize, hh),
                      wnum * cropSize:min((wnum + 1) * cropSize, ww), :]
                # print(sub.shape)
                imgList.append(sub)

    # 图片尺寸若大于切分尺寸则取imgList图片进行超分辨率处理
    if hh > cropSize and ww > cropSize:
        resultList = []
        for im in tqdm.tqdm(imgList):
            # 调用预测接口
            ##若是送入H,W,C的bgr格式的图片用run_image,找遍文档没找到，只能翻源码才找到的api。。。（PS:已向Paddlehub人员反应了）
            res = sr_model.run_image(im)
            resultList.append(res)

    else:  ##本身已是小图则直接进行超分辨率处理
        result = sr_model.run_image(img)

    # 图片尺寸若大于切分尺寸则取imgList图片进行超分辨率处理
    if hh > cropSize and ww > cropSize:
        ##放大比例，为了通用性，通过处理后图片尺寸除原图片尺寸来获取，没有写死为4
        reSizeRatio = int(resultList[0].shape[0] / cropSize)
        # print('reSizeRatio',reSizeRatio)
        ##输出结果，先生成白色底图
        result = np.zeros((int(hh * reSizeRatio), int(ww * reSizeRatio), 3), dtype=np.int)
        newhh, newww = result.shape[0:2]
        # print('result shape',newhh,newww)
        index = 0
        for hnum in range(hCutNum):
            for wnum in range(wCutNum):
                result[int(hnum * cropSize * reSizeRatio):int(min((hnum + 1) * cropSize * reSizeRatio, newhh)),
                int(wnum * cropSize * reSizeRatio):int(min((wnum + 1) * cropSize * reSizeRatio, newww)),
                :] = resultList[index]
                index += 1
        result = np.array(result, dtype=np.int)
    cv2.imwrite('/mnt/data/paddelocr/pubtabnet/pubtabnet/ToImgs/' + os.path.basename(paths).split('.')[0] + '.png', result)

    print('/mnt/data/paddelocr/pubtabnet/pubtabnet/ToImgs/' + os.path.basename(paths).split('.')[0] + '.png')

# 设置图像路径
for i in path_y:
    print(i)
    ChaoFen(i)
    print('done')